package flowdroid.taint;

import db.DatabaseManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import pojo.MyDataFlowResult;
import pojo.MyTuple;
import soot.Scene;
import soot.SootMethod;
import soot.jimple.infoflow.handlers.ResultsAvailableHandler;
import soot.jimple.infoflow.results.InfoflowResults;
import soot.jimple.infoflow.results.ResultSinkInfo;
import soot.jimple.infoflow.results.ResultSourceInfo;
import soot.jimple.infoflow.solver.cfg.IInfoflowCFG;
import soot.util.HashMultiMap;
import soot.util.MultiMap;
import tuple.TupleManager;
import tuple.data.DataKind;
import tuple.entity.MyEntity;

import java.io.*;
import java.util.HashMap;
import java.util.Set;

public class MyInfoflowHandler implements ResultsAvailableHandler {
    private IInfoflowCFG cfg;
    private final TupleManager manager;
    private final Logger logger = LoggerFactory.getLogger(getClass());
    private final String dumpDir;
    private ResultTupleCreator tupleCreator;

    private boolean recordFlowLog = true;

    public MyInfoflowHandler(TupleManager manager, boolean recordFlowLog) {
        this.manager = manager;
        this.dumpDir = manager.dumpDir;
        this.recordFlowLog = recordFlowLog;
    }

    @Override
    public void onResultsAvailable(IInfoflowCFG cfg, InfoflowResults results) {
        this.cfg = cfg;
        System.gc();
        Scene.v().loadNecessaryClasses();
        if (results != null && !results.isEmpty()) {
            MultiMap<ResultSinkInfo, ResultSourceInfo> sinkMapSource = results.getResults();
            int originPathNum = checkOriginResults(sinkMapSource);

            logger.info(manager.phase.toString());
            logger.info("Flow Results Size:" + results.size());
            logger.info("Detected Leak Cnt: " + sinkMapSource.size());
            logger.info("Detected Path Cnt: " + originPathNum);

            try {
                this.tupleCreator = new ResultTupleCreator();
                MultiMap<ResultSinkInfo, ResultSourceInfo> filterResults = filterResults(results);
                processResultsAndCreateTuple(filterResults);
                dumpTupleManagerInfoToFile(dumpDir + File.separator + manager.phase.toString());
            } catch (Exception e) {
                logger.error(e.getMessage());
            }

        }
    }

    private void processResultsAndCreateTuple(MultiMap<ResultSinkInfo, ResultSourceInfo> filterResults) {
        int leakCnt = 0;
        int filteredPathCnt = 0;
        for (ResultSinkInfo sink : filterResults.keySet()) {
            leakCnt++;
            Set<ResultSourceInfo> srcs = filterResults.get(sink);

            for (ResultSourceInfo source : srcs) {
                filteredPathCnt++;
                logger.info(String.format("leak %d, filtered path %d: [%s -> %s]; ",
                        leakCnt, filteredPathCnt, source.toString(), sink.toString()));

                try {
                    MyDataFlowResult result = new MyDataFlowResult(source, sink, cfg);
                    MyTuple myTuple = tupleCreator.apply(result);

                    if (myTuple != null) {
                        logger.info(String.format("Tuple: %s", myTuple.toString()));
                        manager.addTuple(myTuple);
                    }
                } catch (RuntimeException re) {
                    logger.error(re.getMessage());
                }

            }
        }

        manager.syncIndexAfterChange();
        try {
            insertFlowResultsIntoDB(leakCnt, filteredPathCnt);
        } catch (Exception e) {
            logger.error("Insert Tuple Error: " + e.getMessage());
            logger.error("This DB Error is in file: MyInfoflowHandler.java, and caused by invoking insertFlowResultsIntoDB(...)");
        }
    }


    private int checkOriginResults(MultiMap<ResultSinkInfo, ResultSourceInfo> sinkMapSource) {
        int pathCnt = 0;
        for (ResultSinkInfo sink : sinkMapSource.keySet()) {
            Set<ResultSourceInfo> originSrcs = sinkMapSource.get(sink);
            pathCnt += originSrcs.size();
            insertSourceSinkResultIntoDB(sink, originSrcs);
        }
        return pathCnt;
    }

    /**
     * Filters the results to remove redundant paths.
     *
     * @param results The information flow results.
     * @return A MultiMap of filtered results.
     */
    private MultiMap<ResultSinkInfo, ResultSourceInfo> filterResults(InfoflowResults results) {
        ResultFilter filter = new ResultFilter(cfg);

        MultiMap<ResultSinkInfo, ResultSourceInfo> origin = results.getResults();
        MultiMap<ResultSinkInfo, ResultSourceInfo> filtered = new HashMultiMap<>();

        for (ResultSinkInfo sink : origin.keySet()) {
            Set<ResultSourceInfo> originSrcs = origin.get(sink);
            Set<ResultSourceInfo> filteredSources = filter.apply(originSrcs);
            filtered.putAll(sink, filteredSources);
        }
        return filtered;
    }

    private void insertSourceSinkResultIntoDB(ResultSinkInfo sink, Set<ResultSourceInfo> originSrcs) {
        try {
            // SINK
            int freq = originSrcs.size();
            SootMethod sinkSM = sink.getStmt().getInvokeExpr().getMethod();
            DatabaseManager.insertSourceSinkFreq(manager.apkName, sinkSM.getSignature(),
                    "SINK", freq);

            // SOURCES
            HashMap<String, Integer> sourceSigMapFreq = new HashMap<>();
            for (ResultSourceInfo source : originSrcs) {
                String signature = source.getStmt().getInvokeExpr().getMethod().getSignature();
                sourceSigMapFreq.put(signature, sourceSigMapFreq.getOrDefault(signature, 0) + 1);
            }
            for (String sourceSig : sourceSigMapFreq.keySet()) {
                DatabaseManager.insertSourceSinkFreq(manager.apkName, sourceSig,
                        "SOURCE", sourceSigMapFreq.get(sourceSig));
            }
        } catch (Exception e) {
            logger.error("Insert Source Sink Freq Error: " + e.getMessage());
            logger.error("A DB error occurred. In file: MyInfoflowHandler.java, and caused by invoking insertSourceSinkResultIntoDB(...)");
        }
    }

    private void insertFlowResultsIntoDB(int leakCnt, int pathCnt) {
        try {
            Set<DataKind> data = manager.getData();
            Set<MyEntity> entity = manager.getEntity();
            Set<String> permissions = manager.getPermissions();
            int conditionCnt = (int) manager.tuples.stream().filter(t -> t.stringMarks.size() > 0).count();
            DatabaseManager.insertFlowResult(manager.apkName, manager.phase.toString(),
                    manager.tuples.size(), leakCnt, pathCnt,
                    data.size(), entity.size(), conditionCnt,
                    manager.tuples.toString(),
                    data.toString(), entity.toString(), permissions.toString()
            );
        } catch (Exception e) {
            logger.error("Insert Flow Result Error: " + e.getMessage());
            logger.error("A DB error occurred. In file: MyInfoflowHandler.java, and caused by invoking insertFlowResultsIntoDB(...)");
            System.err.println("A DB error occurred. In file: MyInfoflowHandler.java, and caused by invoking insertFlowResultsIntoDB(...)");
        }
    }


    private void dumpTupleManagerInfoToFile(String outputPath) {
        Set<DataKind> data = manager.getData();
        Set<MyEntity> entity = manager.getEntity();
        Set<String> permissions = manager.getPermissions();

        try (BufferedWriter writer = recordFlowLog ?
                new BufferedWriter(new FileWriter(outputPath, true)) :
                new BufferedWriter(new OutputStreamWriter(System.out))
        ) {
            writer.write("size:" + manager.tuples.size());
            writer.newLine();
            writer.write(manager.tuples.toString());
            writer.newLine();
            writer.write("all entity related strings; size:" + entity.size() + ";" + entity);
            writer.newLine();
            writer.write("all data item related strings; size:" + data.size() + ";" + data);
            writer.newLine();
            writer.write("all permission related strings; size:" + permissions.size() + ";" + permissions);
        } catch (IOException e) {
            logger.error("Insert Tuple Error: " + e.getMessage());
        }
    }
}
